import os, sys, json, argparse, subprocess

def command(cmd):
    try:
        rtn = subprocess.check_output(cmd.split(' '))
    except subprocess.CalledProcessError as err:
        print(err)
        rtn = None
    return rtn 

class Downloader(object):
    def __init__(self, exe_path):
        self.dl_formatter = exe_path + " {} -o {}"
        # print(command("which youtube-dl"))

    def download_once(self, vid, output_dir):
        output_path = os.path.join(output_dir, vid+'.mp4')
        if os.path.exists(output_path):
            print(f"file {output_path} exists, skip it.")
            return 
        else:
            print(f"download {vid} to {output_path}")
        rtn = command(self.dl_formatter.format(vid, output_path))
        if rtn is None:
            print(f"download {vid} failed.")
        
    def download_list(self, vid_list, output_dir):
        for idx, vid in enumerate(vid_list):
            print(f"{idx}/{len(vid_list)}")
            self.download_once(vid, output_dir)


def get_args():
    parser = argparse.ArgumentParser("youtube download")
    parser.add_argument("--json_file", type=str, default="./Youtube-Birds/train_list.json")
    parser.add_argument("--output_folder", type=str, default="./output/Youtube-Birds")
    parser.add_argument("--exe_path", type=str, default="D:\\software\youtube-dl\youtube-dl.exe")
    
    args = parser.parse_args()
    return args  


def main():
    args = get_args()
    downloader = Downloader(args.exe_path)
    # read json
    with open(args.json_file, 'r') as fp:
        data = json.load(fp)
    vid_list = [d[0] for d in data]
    output_dir = os.path.abspath(args.output_folder)
    downloader.download_list(vid_list, output_dir)


if __name__=="__main__":
    main()